# -*- coding: utf-8 -*-
# @Time    : 2021/9/28 18:27
# @Author  : huangwei
# @File    : roll_image.py
# @Software: PyCharm
import imutils as imutils
from PIL import Image
from skimage import measure
import math
import cv2

from table import table_predict
from table_method import get_line, draw_lines


def roll_img( filepath, tableNet, output_filepath ):
    """
    用于将图片转正，找到所有长度大于100的横线进行判断
    :param output_filepath:
    :param tableNet:
    :param filepath:
    :return:
    """
    img = Image.open(filepath).convert('RGB')
    prob = 0.2  # 是否连通的概率判断
    row = 100  # 短于这个长度的横线舍去

    # 将其转为正方形再传入模型中进行识别
    width = max(img.size)
    width = (width // 128 + 1) * 128
    input_size = (width, width)

    out, fx, fy = table_predict(img, tableNet, input_size)
    rows = out[0]

    # 实现连通区域标记，connectivity=2 表示8连通，即该点周围8个点都与其值相等，该处标记则为True
    labels = measure.label(rows > prob, connectivity=2)
    # 返回所有连通区域的属性列表
    regions = measure.regionprops(labels)
    # 该连通区域的边界外接框宽度大于阈值row则将该外接框的横向中心线加入到集合中
    row_lines = [get_line(line.coords) for line in regions if line.bbox[3] - line.bbox[1] > row]

    # 查看找到的所有横线
    # 将所有找出的线在黑色图片上画出，找出只有这些线存在的图片上的连通区域
    import numpy as np
    tmp = np.zeros(input_size[::-1], dtype='uint8')
    tmp = draw_lines(tmp, row_lines, color=255, line_width=1)
    cv2.imwrite('images/find_rows.png', tmp)

    # 计算所找到的所有横线的斜率
    k_list = []
    for row_line in row_lines:
        x1, y1, x2, y2 = row_line

        # 不可能为竖线所以 x1! = x2
        if x2 - x1 != 0:
            k = (y2 - y1) / (x2 - x1)
            k_list.append(k)
    k_sort_list = sorted(k_list)

    # 取中位数作为斜率，求旋转角度
    row_length = len(k_sort_list)
    angle = k_sort_list[row_length // 2]
    angle = math.degrees(math.atan(angle))

    # 进行旋转
    img2 = cv2.imread(filepath)
    result_img = imutils.rotate_bound(img2, -angle)
    cv2.imwrite(output_filepath, result_img)


# WEIGHT_PATH = "models/table.weights"
# tableNet = cv2.dnn.readNetFromDarknet(WEIGHT_PATH.replace('.weights', '.cfg'), WEIGHT_PATH)
# img_path = "images/roll2.png"
# output_filepath = "output/roll2_img.png"
# roll_img(img_path, tableNet, output_filepath)
